# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import csv
import sys
import os
import pathlib
from collections import namedtuple

from absl import (flags, app)

from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive
from coin.support.feed_tool.feed_stats.logic.util import *

FLAGS = flags.FLAGS

CsvWriter = namedtuple('CsvWriter', ['writer', 'file'])


class DumpStrategy(object):
  def __init__(self, products, csv_root, trading_date):
    self._products = products
    self._csv_root = os.path.expanduser(csv_root)
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._trading_date = trading_date
    self._csv_writers = {}

  def prepare_csv_file(self, product):
    csv_file = open(os.path.join(
        self._csv_root, '%s_%s.csv' % (product.full_symbol, self._trading_date.strftime('%Y%m%d'))),
                    'w',
                    encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = [
        'timestamp',
        'exchange_timestamp',
        'ask_price',
        'bid_price',
        'mid_price',
        'fair_price',
        'fair_basis_rate',
        'funding_rate',
        'indicative_funding_rate',
        'last_change_percent',
        'last_price',
        'mark_price',
        'open_interest',
        'open_value',
        'total_volume',
        'total_turnover',
        'volume24h',
        'turnover24h'
    ]
    writer.writerow(csv_header)
    csv_file.flush()
    return writer, csv_file

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if book_builder_name == get_subscription_key(product):
        book_builder.subscribe_instrument(product, self.on_instrument)
        exchange, symbol = to_book_builder_exchange_symbol(product)
        writer, csv_file = self.prepare_csv_file(product)
        self._csv_writers[(exchange, symbol)] = CsvWriter(writer, csv_file)

  def on_instrument(self, book):
    writer = self._csv_writers.get((book.exchange, book.native_symbol), None)
    assert writer is not None

    timestamp = book.timestamp
    instrument = book.inst
    csv_line = [
        timestamp,
        instrument['timestamp'],
        instrument['askPrice'],
        instrument['bidPrice'],
        instrument['midPrice'],
        instrument['fairPrice'],
        instrument['fairBasisRate'],
        instrument['fundingRate'],
        instrument['indicativeFundingRate'],
        instrument['lastChangePcnt'],
        instrument['lastPrice'],
        instrument['markPrice'],
        instrument['openInterest'],
        instrument['openValue'],
        instrument['totalVolume'],
        instrument['totalTurnover'],
        instrument['volume24h'],
        instrument['turnover24h']
    ]
    writer.writer.writerow(csv_line)
    writer.file.flush()

  def close_csvs(self):
    for writer in self._csv_writers.values():
      writer.file.flush()
      writer.file.close()


def launch(trading_date, machine, hours, csv_root):
  bitmex_product1 = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')
  bitmex_product2 = BitmexFuturesProduct.FromStr('ETH-USD.PERPETUAL')
  products = [bitmex_product1, bitmex_product2]

  start_time = trading_date  # UTC 0 hour

  strategy = DumpStrategy(products, csv_root, trading_date)

  run_from_archive(products,
                   strategy.on_book_reset,
                   start_time,
                   start_time + datetime.timedelta(hours=hours),
                   machine=[machine])

  strategy.close_csvs()
  return 0


def main(argv):
  trading_date = FLAGS.trading_date
  machine = FLAGS.machine
  csv_root = FLAGS.csv_root
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert csv_root, '--csv_root must be specified.'
  print('Running for %s %s ...' % (trading_date, machine))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch(trading_date, machine, FLAGS.hours, csv_root)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  flags.DEFINE_integer('hours', 24, 'Time span in hours from 00:00:00 UTC.')

  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  app.run(main)
